#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
TimescaleDB数据库模型和操作类

提供GitLab数据的数据库操作接口，包括：
- 数据库连接管理
- 表结构操作
- 数据插入、更新、查询
- 批量操作优化
"""

from contextlib import contextmanager
from datetime import datetime, date
from typing import Dict, List, Optional

import psycopg2
import psycopg2.extras
from loguru import logger
from psycopg2.pool import ThreadedConnectionPool
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine


class TimescaleDBManager:
    """TimescaleDB数据库管理器"""

    def __init__(self, connection_string: str, min_connections: int = 1, max_connections: int = 10):
        """
        初始化数据库管理器
        
        Args:
            connection_string: 数据库连接字符串
            min_connections: 最小连接数
            max_connections: 最大连接数
        """
        self.connection_string = connection_string
        self.pool = None
        self.min_connections = min_connections
        self.max_connections = max_connections
        self._engine = None
        self._initialize_pool()
        self._initialize_engine()

    def _initialize_pool(self):
        """初始化连接池"""
        try:
            self.pool = ThreadedConnectionPool(
                self.min_connections,
                self.max_connections,
                self.connection_string
            )
            logger.info(f"数据库连接池初始化成功，连接数: {self.min_connections}-{self.max_connections}")
        except Exception as e:
            logger.error(f"数据库连接池初始化失败: {e}")
            raise

    def _initialize_engine(self):
        """初始化SQLAlchemy引擎"""
        try:
            self._engine = create_engine(
                self.connection_string,
                pool_size=self.max_connections,
                max_overflow=0,
                pool_pre_ping=True
            )
            logger.info("SQLAlchemy引擎初始化成功")
        except Exception as e:
            logger.error(f"SQLAlchemy引擎初始化失败: {e}")
            raise

    def get_engine(self) -> Engine:
        """获取SQLAlchemy引擎"""
        return self._engine

    @contextmanager
    def get_connection(self):
        """获取数据库连接的上下文管理器"""
        conn = None
        try:
            conn = self.pool.getconn()
            yield conn
        except Exception as e:
            if conn:
                conn.rollback()
            logger.error(f"数据库操作异常: {e}")
            raise
        finally:
            if conn:
                self.pool.putconn(conn)

    def test_connection(self) -> bool:
        """测试数据库连接"""
        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute("SELECT version();")
                    version = cursor.fetchone()[0]
                    logger.info(f"数据库连接成功: {version}")

                    # 检查TimescaleDB扩展
                    cursor.execute("SELECT extname FROM pg_extension WHERE extname = 'timescaledb';")
                    result = cursor.fetchone()
                    if result:
                        logger.info("TimescaleDB扩展已安装")
                    else:
                        logger.warning("TimescaleDB扩展未安装")

                    return True
        except Exception as e:
            logger.error(f"数据库连接测试失败: {e}")
            return False

    def execute_schema_file(self, schema_file_path: str) -> bool:
        """执行SQL架构文件"""
        try:
            with open(schema_file_path, 'r', encoding='utf-8') as f:
                schema_sql = f.read()

            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(schema_sql)
                    conn.commit()
                    logger.info(f"数据库架构文件执行成功: {schema_file_path}")
                    return True
        except Exception as e:
            logger.error(f"执行数据库架构文件失败: {e}")
            return False

    def insert_project(self, project_data: Dict) -> bool:
        """插入或更新项目信息"""
        sql = """
        INSERT INTO gitlab_projects (
            gitlab_url, project_id, project_name, project_path, project_path_with_namespace,
            project_description, default_branch, visibility, web_url, ssh_url_to_repo,
            http_url_to_repo, namespace_name, namespace_path, owner_name, owner_username,
            created_at, updated_at, last_activity_at, star_count, forks_count,
            issues_enabled, merge_requests_enabled, wiki_enabled, archived, empty_repo
        ) VALUES (
            %(gitlab_url)s, %(project_id)s, %(project_name)s, %(project_path)s, %(project_path_with_namespace)s,
            %(project_description)s, %(default_branch)s, %(visibility)s, %(web_url)s, %(ssh_url_to_repo)s,
            %(http_url_to_repo)s, %(namespace_name)s, %(namespace_path)s, %(owner_name)s, %(owner_username)s,
            %(created_at)s, %(updated_at)s, %(last_activity_at)s, %(star_count)s, %(forks_count)s,
            %(issues_enabled)s, %(merge_requests_enabled)s, %(wiki_enabled)s, %(archived)s, %(empty_repo)s
        )
        ON CONFLICT (gitlab_url, project_id) DO UPDATE SET
            project_name = EXCLUDED.project_name,
            project_path = EXCLUDED.project_path,
            project_path_with_namespace = EXCLUDED.project_path_with_namespace,
            project_description = EXCLUDED.project_description,
            default_branch = EXCLUDED.default_branch,
            visibility = EXCLUDED.visibility,
            web_url = EXCLUDED.web_url,
            ssh_url_to_repo = EXCLUDED.ssh_url_to_repo,
            http_url_to_repo = EXCLUDED.http_url_to_repo,
            namespace_name = EXCLUDED.namespace_name,
            namespace_path = EXCLUDED.namespace_path,
            owner_name = EXCLUDED.owner_name,
            owner_username = EXCLUDED.owner_username,
            updated_at = EXCLUDED.updated_at,
            last_activity_at = EXCLUDED.last_activity_at,
            star_count = EXCLUDED.star_count,
            forks_count = EXCLUDED.forks_count,
            issues_enabled = EXCLUDED.issues_enabled,
            merge_requests_enabled = EXCLUDED.merge_requests_enabled,
            wiki_enabled = EXCLUDED.wiki_enabled,
            archived = EXCLUDED.archived,
            empty_repo = EXCLUDED.empty_repo,
            sync_time = NOW()
        """

        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, project_data)
                    conn.commit()
                    return True
        except Exception as e:
            logger.error(f"插入项目数据失败: {e}")
            return False

    def insert_user(self, user_data: Dict) -> bool:
        """插入或更新用户信息"""
        sql = """
        INSERT INTO gitlab_users (
            gitlab_url, user_id, username, name, email, public_email, private_email, state,
            avatar_url, web_url, bio, location, organization, job_title,
            created_at, last_sign_in_at, last_activity_on, projects_limit,
            can_create_group, can_create_project, two_factor_enabled, external, private_profile
        ) VALUES (
            %(gitlab_url)s, %(user_id)s, %(username)s, %(name)s, %(email)s, %(public_email)s, %(private_email)s, %(state)s,
            %(avatar_url)s, %(web_url)s, %(bio)s, %(location)s, %(organization)s, %(job_title)s,
            %(created_at)s, %(last_sign_in_at)s, %(last_activity_on)s, %(projects_limit)s,
            %(can_create_group)s, %(can_create_project)s, %(two_factor_enabled)s, %(external)s, %(private_profile)s
        )
        ON CONFLICT (gitlab_url, user_id) DO UPDATE SET
            username = EXCLUDED.username,
            name = EXCLUDED.name,
            email = EXCLUDED.email,
            public_email = EXCLUDED.public_email,
            private_email = EXCLUDED.private_email,
            state = EXCLUDED.state,
            avatar_url = EXCLUDED.avatar_url,
            web_url = EXCLUDED.web_url,
            bio = EXCLUDED.bio,
            location = EXCLUDED.location,
            organization = EXCLUDED.organization,
            job_title = EXCLUDED.job_title,
            last_sign_in_at = EXCLUDED.last_sign_in_at,
            last_activity_on = EXCLUDED.last_activity_on,
            projects_limit = EXCLUDED.projects_limit,
            can_create_group = EXCLUDED.can_create_group,
            can_create_project = EXCLUDED.can_create_project,
            two_factor_enabled = EXCLUDED.two_factor_enabled,
            external = EXCLUDED.external,
            private_profile = EXCLUDED.private_profile,
            sync_time = NOW()
        """

        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, user_data)
                    conn.commit()
                    return True
        except Exception as e:
            logger.error(f"插入用户数据失败: {e}")
            return False

    def insert_branch(self, branch_data: Dict) -> bool:
        """插入或更新分支信息"""
        sql = """
        INSERT INTO gitlab_branches (
            gitlab_url, project_id, branch_name, branch_commit_sha, branch_commit_title,
            branch_commit_message, branch_commit_author_name, branch_commit_author_email,
            branch_commit_date, is_default, is_protected, can_push,
            developers_can_push, developers_can_merge
        ) VALUES (
            %(gitlab_url)s, %(project_id)s, %(branch_name)s, %(branch_commit_sha)s, %(branch_commit_title)s,
            %(branch_commit_message)s, %(branch_commit_author_name)s, %(branch_commit_author_email)s,
            %(branch_commit_date)s, %(is_default)s, %(is_protected)s, %(can_push)s,
            %(developers_can_push)s, %(developers_can_merge)s
        )
        ON CONFLICT (gitlab_url, project_id, branch_name) DO UPDATE SET
            branch_commit_sha = EXCLUDED.branch_commit_sha,
            branch_commit_title = EXCLUDED.branch_commit_title,
            branch_commit_message = EXCLUDED.branch_commit_message,
            branch_commit_author_name = EXCLUDED.branch_commit_author_name,
            branch_commit_author_email = EXCLUDED.branch_commit_author_email,
            branch_commit_date = EXCLUDED.branch_commit_date,
            is_default = EXCLUDED.is_default,
            is_protected = EXCLUDED.is_protected,
            can_push = EXCLUDED.can_push,
            developers_can_push = EXCLUDED.developers_can_push,
            developers_can_merge = EXCLUDED.developers_can_merge,
            sync_time = NOW()
        """

        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, branch_data)
                    conn.commit()
                    return True
        except Exception as e:
            logger.error(f"插入分支数据失败: {e}")
            return False

    def insert_commit(self, commit_data: Dict) -> bool:
        """插入提交数据"""
        sql = """
        INSERT INTO gitlab_commits (
            gitlab_url, project_branch_key, commit_time, commit_sha, project_id, project_name,
            project_path_with_namespace, branch_name, commit_title, commit_message,
            commit_short_id, author_name, author_email, author_username, author_user_id,
            committer_name, committer_email, committer_username, committer_user_id,
            additions, deletions, total_changes, files_changed, commit_type,
            is_merge_commit, parent_ids, authored_date, committed_date
        ) VALUES (
            %(gitlab_url)s, %(project_branch_key)s, %(commit_time)s, %(commit_sha)s, %(project_id)s, %(project_name)s,
            %(project_path_with_namespace)s, %(branch_name)s, %(commit_title)s, %(commit_message)s,
            %(commit_short_id)s, %(author_name)s, %(author_email)s, %(author_username)s, %(author_user_id)s,
            %(committer_name)s, %(committer_email)s, %(committer_username)s, %(committer_user_id)s,
            %(additions)s, %(deletions)s, %(total_changes)s, %(files_changed)s, %(commit_type)s,
            %(is_merge_commit)s, %(parent_ids)s, %(authored_date)s, %(committed_date)s
        )
        ON CONFLICT (gitlab_url, project_branch_key, commit_time, commit_sha) DO NOTHING
        """

        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, commit_data)
                    conn.commit()
                    return True
        except Exception as e:
            logger.error(f"插入提交数据失败: {e}")
            return False

    def batch_insert_commits(self, commits_data: List[Dict]) -> int:
        """批量插入提交数据"""
        if not commits_data:
            return 0

        sql = """
        INSERT INTO gitlab_commits (
            gitlab_url, project_branch_key, commit_time, commit_sha, project_id, project_name,
            project_path_with_namespace, branch_name, commit_title, commit_message,
            commit_short_id, author_name, author_email, author_username, author_user_id,
            committer_name, committer_email, committer_username, committer_user_id,
            additions, deletions, total_changes, files_changed, commit_type,
            is_merge_commit, parent_ids, authored_date, committed_date
        ) VALUES (
            %(gitlab_url)s, %(project_branch_key)s, %(commit_time)s, %(commit_sha)s, %(project_id)s, %(project_name)s,
            %(project_path_with_namespace)s, %(branch_name)s, %(commit_title)s, %(commit_message)s,
            %(commit_short_id)s, %(author_name)s, %(author_email)s, %(author_username)s, %(author_user_id)s,
            %(committer_name)s, %(committer_email)s, %(committer_username)s, %(committer_user_id)s,
            %(additions)s, %(deletions)s, %(total_changes)s, %(files_changed)s, %(commit_type)s,
            %(is_merge_commit)s, %(parent_ids)s, %(authored_date)s, %(committed_date)s
        )
        ON CONFLICT (gitlab_url, project_branch_key, commit_time, commit_sha) DO NOTHING
        """

        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.executemany(sql, commits_data)
                    affected_rows = cursor.rowcount
                    conn.commit()
                    logger.info(f"批量插入提交数据成功，影响行数: {affected_rows}")
                    return affected_rows
        except Exception as e:
            logger.error(f"批量插入提交数据失败: {e}")
            return 0

    def get_sync_status(self, gitlab_url: str, project_id: int, branch_name: str) -> Optional[Dict]:
        """获取同步状态"""
        sql = """
        SELECT gitlab_url, project_id, branch_name, last_sync_time, last_commit_time, 
               last_commit_sha, sync_type, sync_status, total_commits_synced
        FROM gitlab_sync_status 
        WHERE gitlab_url = %s AND project_id = %s AND branch_name = %s
        """

        try:
            with self.get_connection() as conn:
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    cursor.execute(sql, (gitlab_url, project_id, branch_name))
                    result = cursor.fetchone()
                    return dict(result) if result else None
        except Exception as e:
            logger.error(f"获取同步状态失败: {e}")
            return None

    def update_sync_status(self, gitlab_url: str, project_id: int, branch_name: str,
                           last_commit_time: Optional[datetime] = None,
                           last_commit_sha: Optional[str] = None,
                           sync_status: str = 'completed',
                           total_commits_synced: int = 0) -> bool:
        """更新同步状态"""
        sql = """
        INSERT INTO gitlab_sync_status (
            gitlab_url, project_id, branch_name, last_sync_time, last_commit_time,
            last_commit_sha, sync_status, total_commits_synced, updated_at
        ) VALUES (
            %s, %s, %s, NOW(), %s, %s, %s, %s, NOW()
        )
        ON CONFLICT (gitlab_url, project_id, branch_name) DO UPDATE SET
            last_sync_time = NOW(),
            last_commit_time = EXCLUDED.last_commit_time,
            last_commit_sha = EXCLUDED.last_commit_sha,
            sync_status = EXCLUDED.sync_status,
            total_commits_synced = gitlab_sync_status.total_commits_synced + EXCLUDED.total_commits_synced,
            updated_at = NOW()
        """

        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, (gitlab_url, project_id, branch_name, last_commit_time,
                                         last_commit_sha, sync_status, total_commits_synced))
                    conn.commit()
                    return True
        except Exception as e:
            logger.error(f"更新同步状态失败: {e}")
            return False

    def get_project_branches_stats(self, gitlab_url: Optional[str] = None, project_id: Optional[int] = None) -> List[
        Dict]:
        """获取项目分支统计"""
        where_conditions = []
        params = []

        if gitlab_url:
            where_conditions.append("gitlab_url = %s")
            params.append(gitlab_url)

        if project_id:
            where_conditions.append("project_id = %s")
            params.append(project_id)

        where_clause = "WHERE " + " AND ".join(where_conditions) if where_conditions else ""

        sql = f"""
        SELECT * FROM v_project_branch_stats
        {where_clause}
        ORDER BY total_commits DESC
        """

        try:
            with self.get_connection() as conn:
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    cursor.execute(sql, params)
                    results = cursor.fetchall()
                    return [dict(row) for row in results]
        except Exception as e:
            logger.error(f"获取项目分支统计失败: {e}")
            return []

    def get_user_contribution_stats(self, gitlab_url: Optional[str] = None,
                                    start_date: Optional[date] = None,
                                    end_date: Optional[date] = None) -> List[Dict]:
        """获取用户贡献统计"""
        where_conditions = []
        params = []

        # if gitlab_url:
        #     where_conditions.append("gitlab_url = %s")
        #     params.append(gitlab_url)

        if start_date:
            where_conditions.append("first_commit_time >= %s")
            params.append(start_date)

        if end_date:
            where_conditions.append("last_commit_time <= %s")
            params.append(end_date)

        where_clause = "WHERE " + " AND ".join(where_conditions) if where_conditions else ""

        sql = f"""
        SELECT * FROM v_user_contribution_stats
        {where_clause}
        ORDER BY total_commits DESC
        """
        logger.info(sql)

        try:
            with self.get_connection() as conn:
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    cursor.execute(sql, params)
                    results = cursor.fetchall()
                    return [dict(row) for row in results]
        except Exception as e:
            logger.error(f"获取用户贡献统计失败: {e}")
            return []

    def get_daily_activity_stats(self, gitlab_url: Optional[str] = None,
                                 start_date: Optional[date] = None,
                                 end_date: Optional[date] = None) -> List[Dict]:
        """获取每日活跃度统计"""
        where_conditions = []
        params = []

        if gitlab_url:
            where_conditions.append("gitlab_url = %s")
            params.append(gitlab_url)

        if start_date:
            where_conditions.append("commit_date >= %s")
            params.append(start_date)

        if end_date:
            where_conditions.append("commit_date <= %s")
            params.append(end_date)

        where_clause = "WHERE " + " AND ".join(where_conditions) if where_conditions else ""

        sql = f"""
        SELECT * FROM v_daily_activity_stats
        {where_clause}
        ORDER BY commit_date DESC
        """

        try:
            with self.get_connection() as conn:
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    cursor.execute(sql, params)
                    results = cursor.fetchall()
                    return [dict(row) for row in results]
        except Exception as e:
            logger.error(f"获取每日活跃度统计失败: {e}")
            return []

    def refresh_materialized_views(self) -> bool:
        """刷新用户贡献统计物化视图"""
        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute("SELECT refresh_user_contribution_stats();")
                    conn.commit()
                    logger.info("物化视图刷新成功")
                    return True
        except Exception as e:
            logger.error(f"刷新物化视图失败: {e}")
            return False

    def close(self):
        """关闭数据库连接池和SQLAlchemy引擎"""
        if self.pool:
            self.pool.closeall()
            logger.info("数据库连接池已关闭")
        if self._engine:
            self._engine.dispose()
            logger.info("SQLAlchemy引擎已关闭")
